# 1D Hermite程序
import numpy as np
from scipy.integrate import quad
from Hermite import *


# 定义 Hermite 基函数及其导数
def phi_0(x):
    return 1 - 3 * x**2 + 2 * x**3


def phi_1(x):
    return x - 2 * x**2 + x**3


def phi_2(x):
    return 3 * x**2 - 2 * x**3


def phi_3(x):
    return -x**2 + x**3


def dphi_0(x):
    return -6 * x + 6 * x**2


def dphi_1(x):
    return 1 - 4 * x + 3 * x**2


def dphi_2(x):
    return 6 * x - 6 * x**2


def dphi_3(x):
    return -2 * x + 3 * x**2


def ddphi_0(x):
    return -6 + 12 * x  # phi_0 的二阶导数


def ddphi_1(x):
    return -4 + 6 * x   # phi_1 的二阶导数


def ddphi_2(x):
    return 6 - 12 * x   # phi_2 的二阶导数


def ddphi_3(x):
    return -2 + 6 * x   # phi_3 的二阶导数



class Mesh:
    def __init__(self, n: int, interval: list = [0, 1]):
        """n为剖分设置
            interval为区间
        Args:
            n (int): _description_
        """
        mesh = np.linspace(interval[0], interval[1], n+1)
        self.mesh = mesh[:]
        self.h = (interval[1] - interval[0]) / n  # 步长
        self.n = n  # 网格剖分
        self.n_nodes = n + 1  # 节点数
        
        
    

class FE:
    def __init__(self, mesh: Mesh, order: int=3):
        """Hermite有限元空间

        Args:
            mesh (Mesh): 网格
            order (int): 有限元空间的阶数
        """
        self.mesh = mesh.mesh  # 网格
        self.n = mesh.n  # 网格剖分
        self.h = mesh.h  # 网格步长
        self.n_nodes = mesh.n_nodes  # 节点数
        self.n_dofs_per_node = 2
        self.n_dofs = self.n_dofs_per_node * mesh.n_nodes  # 自由度个数=节点数*2
        # 初始化刚度矩阵 和 载荷向量
        self.K = np.zeros((self.n_dofs, self.n_dofs))
        self.b = np.zeros(self.n_dofs)
        
    def assemble_matrix(self):
        n_dofs_per_element = 4  # 每个单元4个自由度
        for i in range(self.n):
            Ke = np.zeros((n_dofs_per_element, n_dofs_per_element))
            for alpha in range(n_dofs_per_element):
                # 选择二阶导数函数
                if alpha == 0:
                    dd_alpha = ddphi_0
                elif alpha == 1:
                    dd_alpha = ddphi_1
                elif alpha == 2:
                    dd_alpha = ddphi_2
                else:
                    dd_alpha = ddphi_3

                for beta in range(n_dofs_per_element):
                    if beta == 0:
                        dd_beta = ddphi_0
                    elif beta == 1:
                        dd_beta = ddphi_1
                    elif beta == 2:
                        dd_beta = ddphi_2
                    else:
                        dd_beta = ddphi_3

                    # 积分项调整为二阶导数乘积及正确缩放因子
                    def integrand(x): return (1 / self.h**3) * \
                        dd_alpha(x) * dd_beta(x)
                    Ke[alpha, beta] = quad(integrand, 0, 1)[0]

            # 组装到全局矩阵
            dof_indices = [2*i, 2*i+1, 2*(i+1), 2*(i+1)+1]
            for a in range(4):
                for b in range(4):
                    self.K[dof_indices[a], dof_indices[b]] += Ke[a, b]
                    
                    
    def assemble_vector(self, f):
        """
        组装载荷向量
        """
        n_dofs_per_element = self.n_dofs_per_node * 2  # 每个单元的自由度数量

        for i in range(self.n):
            # 单元载荷向量
            be = np.zeros(n_dofs_per_element)
            for alpha in range(n_dofs_per_element):
                if alpha == 0:
                    phi_alpha = phi_0
                elif alpha == 1:
                    phi_alpha = phi_1
                elif alpha == 2:
                    phi_alpha = phi_2
                else:
                    phi_alpha = phi_3

                def integrand(x):
                    return self.h * f(self.mesh[i] + x * self.h) * phi_alpha(x)

                be[alpha] = quad(integrand, 0, 1)[0]

            # 局部到全局的映射
            dof_indices = [self.n_dofs_per_node * i +
                           j for j in range(n_dofs_per_element)]
            for alpha in range(n_dofs_per_element):
                self.b[dof_indices[alpha]] += be[alpha]
                
        
    def dirichlet_bc(self, left_value, right_value):
        """
        应用狄利克雷边界条件
        :param left_value: 左边界的值
        :param right_value: 右边界的值
        """
        # 左边界自由度索引
        left_dof_index = 0
        # 右边界自由度索引
        right_dof_index = self.n_dofs - 2

        # 修改刚度矩阵
        self.K[left_dof_index, :] = 0
        self.K[left_dof_index, left_dof_index] = 1
        self.K[right_dof_index, :] = 0
        self.K[right_dof_index, right_dof_index] = 1

        # 修改载荷向量
        self.b[left_dof_index] = left_value
        self.b[right_dof_index] = right_value
        
    
    def neumann_bc(self):
        """
        应用诺伊曼边界条件 u'(0) = u'(1) = 0
        """
        # 左边界导数自由度索引
        left_derivative_dof_index = 1
        # 右边界导数自由度索引
        right_derivative_dof_index = self.n_dofs - 1

        # 修改刚度矩阵
        self.K[left_derivative_dof_index, :] = 0
        self.K[left_derivative_dof_index, left_derivative_dof_index] = 1
        self.K[right_derivative_dof_index, :] = 0
        self.K[right_derivative_dof_index, right_derivative_dof_index] = 1

        # 修改载荷向量
        self.b[left_derivative_dof_index] = 0
        self.b[right_derivative_dof_index] = 0
    
    
    def solve(self, f, left_value, right_value):
        """
        求解线性方程组
        """
        self.assemble_matrix()
        self.assemble_vector(f)
        self.dirichlet_bc(left_value, right_value)
        self.neumann_bc()
        self.u = np.linalg.solve(self.K, self.b)
        return self.u
    
    def compute_l2_norm(self, u_coeff):
        """计算数值解的L2范数"""
        l2_norm = 0.0
        for i in range(self.n):
            # 获取当前单元的自由度索引
            dof_indices = [2*i, 2*i+1, 2*(i+1), 2*(i+1)+1]
            coeff = u_coeff[dof_indices]  # [u_i, u'_i, u_{i+1}, u'_{i+1}]

            # 在参考单元[0,1]上积分
            def integrand(x):
                # 数值解表达式: sum(coeff[k] * phi_k(x))
                u_h = coeff[0]*phi_0(x) + coeff[1]*phi_1(x) + \
                    coeff[2]*phi_2(x) + coeff[3]*phi_3(x)
                return u_h**2

            integral = quad(integrand, 0, 1)[0]
            l2_norm += self.h * integral  # 积分变量变换带来的系数h
        return np.sqrt(l2_norm)

    def compute_gradient_l2_norm(self, u_coeff):
        """计算数值解一阶导数的L2范数"""
        grad_norm = 0.0
        for i in range(self.n):
            dof_indices = [2*i, 2*i+1, 2*(i+1), 2*(i+1)+1]
            coeff = u_coeff[dof_indices]

            def integrand(x):
                # 一阶导数表达式: sum(coeff[k] * dphi_k(x)) / h
                du_h = (coeff[0]*dphi_0(x) + coeff[1]*dphi_1(x)
                        + coeff[2]*dphi_2(x) + coeff[3]*dphi_3(x)) / self.h
                return du_h**2

            integral = quad(integrand, 0, 1)[0]
            grad_norm += self.h * integral  # h来自积分变量变换
        return np.sqrt(grad_norm)

    def compute_hessian_l2_norm(self, u_coeff):
        """计算数值解二阶导数的L2范数"""
        hessian_norm = 0.0
        for i in range(self.n):
            dof_indices = [2*i, 2*i+1, 2*(i+1), 2*(i+1)+1]
            coeff = u_coeff[dof_indices]

            def integrand(x):
                # 二阶导数表达式: sum(coeff[k] * ddphi_k(x)) / h^2
                d2u_h = (coeff[0]*ddphi_0(x) + coeff[1]*ddphi_1(x)
                         + coeff[2]*ddphi_2(x) + coeff[3]*ddphi_3(x)) / (self.h**2)
                return d2u_h**2

            integral = quad(integrand, 0, 1)[0]
            hessian_norm += self.h * integral  # h来自积分变量变换
        return np.sqrt(hessian_norm)

    def compute_l2_error(self, u_coeff, true_u):
        """计算与真解的L2误差范数"""
        error = 0.0
        for i in range(self.n):
            x_left = self.mesh[i]
            x_right = self.mesh[i+1]
            dof_indices = [2*i, 2*i+1, 2*(i+1), 2*(i+1)+1]
            coeff = u_coeff[dof_indices]

            def integrand(x):
                # 将参考单元坐标转换为物理坐标
                x_phy = x_left + x * self.h
                # 数值解
                u_h = coeff[0]*phi_0(x) + coeff[1]*phi_1(x) + \
                    coeff[2]*phi_2(x) + coeff[3]*phi_3(x)
                # 真解
                u_true = true_u(x_phy)
                return (u_true - u_h)**2

            integral = quad(integrand, 0, 1)[0]
            error += self.h * integral
        return np.sqrt(error)


